import os
import importlib
import pickle
import json
import torch
import plotly.graph_objs as go
from kornia.color import rgb_to_ycbcr,ycbcr_to_rgb
from ignite.metrics import SSIM 
from model import SDNet as net
from MyDataLoader import VISIRData

gpu=2
if torch.cuda.is_available():
        torch.cuda.set_device(gpu)
        device = 'cuda'
else:
        print('WARNING: [CUDA unavailable] Using CPU instead!')
        device = 'cpu'


# 假设 'model_path' 是保存模型状态的文件路径
model_path = '/home/ubuntu/workspace/SDNet-main/modelsave/SPECT/80/MyModel.pth'  # 替换为实际的文件路径

# 加载保存的状态字典
state = torch.load(model_path)

# 创建模型实例
model = net()  # 假设 DenseNet 是您的模型类

# 将加载的模型参数应用到模型上
model.load_state_dict(state)
model.to(device)
import torchvision.transforms as transforms
# -------------------加载数据集
'''vis_image = torch.cat((vis_image_raw,vis_sobel,vis_laplace,vis_image_equalized),dim=0)'''
dataset_road=VISIRData(set_prefix='/home/ubuntu/workspace/data/roadscene',split='test')
#dataset_m3fd=AuxDatasetTest(data_pth='../data/M3FD_Fusion')
#dataset_road=AuxDatasetTest(data_pth='../data/roadscene')
#dataset_tno=AuxDatasetTest(data_pth='../data/tno')

set_dic={
    'ROAD': dataset_road,
    }

print('device',device)
from tqdm import trange
import torch
import numpy as np
from torchvision import transforms
import skimage as ski
from Evaluator import Evaluator
toGray=transforms.Grayscale()


def metrics(set_name: str,set_dic: dict,model):
    dataset=set_dic.get(set_name)
    len_set=len(dataset)
    with torch.no_grad():
        metric_result = np.zeros((8))
        for i in trange(len_set):
            vis_image, ir_image,refedge=dataset[i]
            fused_img, x1_de, x2_de,edge_out=model(vis_image.unsqueeze(dim=0).to(device),ir_image.unsqueeze(dim=0).to(device))
            vis_image=vis_image.cpu()
            ir_image=ir_image.cpu()
            int_out=fused_img.cpu()
            edge_out=edge_out.cpu()
            #print('vis min',torch.max(ir_image))
            
            joint_in_rvEdge = ((1 - int_out)* edge_out)
            joint_in_noEdge = int_out * (1 - edge_out)
            #fus_edge=joint_in_noEdge+joint_in_rvEdge
            fus_edge=torch.clamp(joint_in_noEdge+joint_in_rvEdge,0,1 )
            '''fus_edge = joint_in_noEdge + joint_in_rvEdge

            # 计算最小值和最大值


            
            min_val = torch.min(fus_edge)
            max_val = torch.max(fus_edge)

            # 重新缩放到 [0, 1] 范围
            fus_edge = (fus_edge - min_val) / (max_val - min_val)'''
            #print(int_out.max())
            '''vi=ski.util.img_as_ubyte(toGray(vis_image).squeeze().numpy()).astype('float32')
            ir=ski.util.img_as_ubyte(ir_image.squeeze().numpy()).astype('float32')
            fi=ski.util.img_as_ubyte(int_out.squeeze().numpy()).astype('float32')'''
            vi=ski.util.img_as_ubyte(toGray(vis_image).squeeze().numpy()).astype('float32')
            ir=ski.util.img_as_ubyte(ir_image.squeeze().numpy()).astype('float32')
            fi=ski.util.img_as_ubyte(fus_edge.squeeze().numpy()).astype('float32')
            metric_result += np.array([Evaluator.EN(fi), Evaluator.SD(fi)
                                                , Evaluator.SF(fi), Evaluator.MI(fi, ir, vi)
                                                , Evaluator.SCD(fi, ir, vi), Evaluator.VIFF(fi, ir, vi)
                                                , Evaluator.Qabf(fi, ir, vi), Evaluator.SSIM(fi, ir, vi)])
            
            '''vi=np.round(cv2.cvtColor(np.transpose(vis_image[:3,:,:].numpy(), (1, 2, 0)), cv2.COLOR_BGR2GRAY)*255)
            ir=np.round(ir_image[0,:,:].squeeze().numpy()*255)
            fi=np.round(fus_edge.squeeze().numpy()*255)
            metric_result += np.array([Evaluator.EN(fi), Evaluator.SD(fi)
                                                , Evaluator.SF(fi), Evaluator.MI(fi, ir, vi)
                                                , Evaluator.SCD(fi, ir, vi), Evaluator.VIFF(fi, ir, vi)
                                                , Evaluator.Qabf(fi, ir, vi), Evaluator.SSIM(fi, ir, vi)])'''
        metric_result /= len_set
    print('Dataset: ',set_name)
    print('\t'+'EN: '+str(np.round(metric_result[0], 2))
                    +', SD: '+str(np.round(metric_result[1], 2))
                    +', SF: '+str(np.round(metric_result[2], 2))
                    +', MI: '+str(np.round(metric_result[3], 2))
                    +', SCD: '+str(np.round(metric_result[4], 2))
                    +', VIFF: '+str(np.round(metric_result[5], 2))
                    +', Qabf: '+str(np.round(metric_result[6], 2))
                    +', SSIM: '+str(np.round(metric_result[7], 2))
                    )
    return None

metrics(set_name='ROAD',set_dic=set_dic,model=model)